import fitz
import os
import random
import string
import argparse
from datetime import datetime

def generate_admin_password(length=16):
    """生成随机管理员密码，包含字母、数字和特殊字符"""
    characters = string.ascii_letters + string.digits + "!@#$%^&*()_+-=[]{}|;:,.<>?"
    return ''.join(random.choice(characters) for _ in range(length))

def add_watermark_encrypt(input_pdf, output_pdf, watermark_text, user_password, admin_password):
    """添加水印并加密PDF，设置不同权限级别"""
    try:
        # 打开PDF文件
        doc = fitz.open(input_pdf)
    
        # 添加水印到每一页
        for page in doc:
            # 获取页面尺寸
            rect = page.rect
            # 创建透明水印文本
            text_width = fitz.get_text_length(watermark_text, fontname="helv", fontsize=12)
            # 计算水印位置 - 对角线放置多个水印
            for i in range(0, int(rect.height), 100):
                for j in range(0, int(rect.width), int(text_width) + 50):
                    # 在当前位置绘制水印
                    page.insert_text(
                        (j, i), 
                        watermark_text, 
                        fontsize=12, 
                        fontname="helv", 
                        rotate=90,  # 使用支持的角度：0, 90, 180 或 270
                        color=(0.7, 0.7, 0.7, 0.3)  # 浅灰色半透明
                    )
        
        # 定义权限 - 使用最小权限（仅阅读）
        user_permissions = 0  # 没有任何特殊权限
        
        try:
            # 尝试使用较新版本的加密方法
            doc.save(
                output_pdf,
                encryption=fitz.PDF_ENCRYPT_AES_256,  # 使用AES 256加密
                user_pw=user_password,        # 用户密码
                owner_pw=admin_password,      # 管理员密码
                permissions=user_permissions  # 用户权限
            )
        except TypeError:
            # 如果上面的方法失败，尝试使用较老版本的加密API
            doc.save(
                output_pdf,
                encryption=fitz.PDF_ENCRYPT_AES_256,  # 使用AES 256加密
                user_pw=user_password,        # 用户密码
                owner_pw=admin_password,      # 管理员密码
                permission=user_permissions   # 注意：旧版本可能使用"permission"而不是"permissions"
            )
        
        # 关闭文档
        doc.close()
        
        return admin_password
    except Exception as e:
        print(f"处理PDF时出错: {e}")
        raise

def main():
    parser = argparse.ArgumentParser(description="给PDF添加水印并加密，设置不同权限级别")
    parser.add_argument("--input", dest="input_pdf", default="input.pdf", 
                  help="输入PDF文件路径 (默认: 当前目录下的input.pdf)")
    parser.add_argument("--output", help="输出PDF文件路径")
    parser.add_argument("--watermark", default="机密文件 - 仅供授权查看", help="水印文本")
    parser.add_argument("--user-pw", default="N!0_123", help="用户密码(只读权限)")
    args = parser.parse_args()

    # 检查输入文件是否存在
    if not os.path.exists(args.input_pdf):
        print(f"错误: 输入文件 '{args.input_pdf}' 不存在")
        return
    
    # 如果没有指定输出文件，则在原文件名后添加后缀
    if not args.output:
        filename, ext = os.path.splitext(args.input_pdf)
        args.output = f"{filename}_watermarked{ext}"
    
    # 生成管理员密码
    admin_password = generate_admin_password()
    
    # 添加当前时间到水印
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M")
    watermark_text = f"{args.watermark} {timestamp}"
    
    # 添加水印并加密
    admin_password = add_watermark_encrypt(
        args.input_pdf,
        args.output,
        watermark_text,
        args.user_pw,
        admin_password
    )
    
    print(f"处理完成！文件已保存为: {args.output}")
    print(f"用户密码 (只读权限): {args.user_pw}")
    print(f"管理员密码 (完全权限): {admin_password}")
    print("请妥善保管管理员密码，它不会被再次显示！")

if __name__ == "__main__":
    main()